import torch
from linear_regression.simple_data_builder import data_iter, synthetic_data

def linreg(X, w, b):
    """
    定义模型
    将模型的输入和参数同模型的输出关联起来
    此案例中
    X是一个1000*2的矩阵
    w是一个2*1的向量
    所以X,w矩阵乘法得到一个(1000,1)的向量
    向量 + 标量(b)，根据广播机制，b将加到向量的每一项之中
    """
    return torch.matmul(X, w) + b


def squared_loss(y_hat, y):
    """
    定义损失函数：均方损失
    y为何要转换为y_hat结构？
    """
    return (y_hat - y.reshape(y_hat.shape))**2 / 2


def sgd(params, lr, batch_size):
    """小批量随机梯度下降
    @batch_size: 批量大小，规范化步长
    @lr: 学习速率，决定每一步更新的大小
    @params: 需要更新的内容，比如权重w，偏移量b
    """
    with torch.no_grad():
        for param in params:
            # param.grad 就是通过backward()触发计算得到的梯度总和
            # param.grad / batch_size 就是平均梯度，保证更新步长与批量大小无关
            # param -= lr * ...  是核心的梯“梯度下降”逻辑，让参数朝着梯度的反方向移动
            param -= lr * param.grad / batch_size
            # 清理梯度，准备下一轮训练
            param.grad.zero_()


if __name__ == '__main__':
    # 通過从均值為0，标准差为0.01的正太分布中采样，初始化权重
    w = torch.normal(0, 0.01, size=(2, 1), requires_grad=True)
    # 将偏置初始化为0
    b = torch.zeros(1, requires_grad=True)
    
    lr = 0.03
    num_epochs = 3  
    net = linreg
    loss = squared_loss

    # 创建虚拟的真实数据
    batch_size = 10
    true_w = torch.tensor([2, -3.4])
    true_b = 4.2
    features, labels = synthetic_data(true_w, true_b, 1000)

    # 训练模型
    for epoch in range(num_epochs):
        for X, y in data_iter(batch_size, features, labels):
            l = loss(net(X, w, b), y)
            l.sum().backward()
            sgd([w,b], lr, batch_size)
        with torch.no_grad():
            train_l = loss(net(features, w, b), labels)
            print(f'epoch {epoch + 1}, loss {float(train_l.mean()):f}')